#!/bin/env python

import numpy as np
import constants

def calc_density_profile( rr, mass_profile, cumulative=1 ) :
    '''
    This function takes in two arrays: rr, a mass_profile, and an optional cumulative argument (default is that the mass profile is cumulative).  
    It returns the corresponding density profile in another array.
    Note: rr, mass_profile, and the resulting density profile must all be arrays of the same length
    Ex:
    hbmpro=load_profiles.load_bmpro_profile(h_bmpro_a?.????.dat,[1],99) 
    icm_den_pro=calc_density_profile(hbmpro['rr'],hbmpro['Mg']) # Returns gas mass profile
    '''
    shell_volumes = 4./3.*np.pi*( rr**3 - np.concatenate( ([0], rr[:len(rr)-1]),axis=0 )**3 )
    if ( cumulative ) :
        shell_masses = mass_profile - np.concatenate( ([0],mass_profile[:len(mass_profile)-1]),axis=0 )
    else :
        shell_masses = mass_profile
    
    return shell_masses/shell_volumes
           
def Ez( hblistr500, aexp ) :
    '''
    Takes in an hblist dictionary for a cluster and the corresponding aexp
    Returns a dictionary of dictionaries with E(z) normalizations:
    Ez = { 'Ez': {}, 'T500': {},
           'P500':{}, 'K500':{}
        }
    Ex:
    hblist = load_hblist( h_blist_r500c_a?.????.dat, clusters=[1,2] )
    Ez = Ez( hblist, ?.???? ) 
    T500 = Ez['T500'][1] # Returns the T500 normalization of halo 1
    '''
    Ez = { 'Ez': {}, 'T500': {},
           'P500':{}, 'K500':{}
        }
    for id in hblistr500['Mtotal'].keys() :
        M500 = hblistr500['Mtotal'][id]
    
        Ez['Ez'][id] = np.sqrt( constants.OmM * (1./float(aexp))**3 + constants.OmL )
        Ez['T500'][id] = 11050. * (M500/1e15)**0.6667 * Ez['Ez'][id]**0.6667 / constants.kB
        Ez['P500'][id] = 1.45 * 1e-11 * (M500/1e15)**0.6667 * Ez['Ez'][id]**2.6667
        Ez['K500'][id] = 1963. * (M500/1e15)**0.6667 * Ez['Ez'][id]**-0.6667

    return Ez
    
